"""
ConnLLM令牌计算模块
提供token计数和从API响应中提取token信息的功能

依赖:
- tiktoken (OpenAI模型)
- anthropic (Anthropic模型)
这些依赖是可选的，如果没有安装则会使用近似估算
"""
import re
import logging
from typing import Dict, Any, Tuple, Optional, List, Union, Set

# 设置日志
logger = logging.getLogger("connllm.tokenizers")

# 记录已经发出的警告，避免重复提示
_warned_missing_libs: Set[str] = set()


def count_tokens(text: str, provider_type: str, model: Optional[str] = None) -> int:
	"""
	计算文本中的token数量
	
	Args:
		text: 要计算的文本
		provider_type: 提供商类型
		model: 模型名称，可选
		
	Returns:
		token数量
	"""
	# 根据提供商类型选择不同的计算方法
	if provider_type == "openai":
		return _count_openai_tokens(text, model)
	elif provider_type == "anthropic":
		return _count_anthropic_tokens(text, model)
	elif provider_type == "openrouter":
		# OpenRouter根据模型前缀选择计算方法
		if model and model.startswith("anthropic/"):
			return _count_anthropic_tokens(text, model.split("/", 1)[1])
		else:
			return _count_openai_tokens(text, model)
	elif provider_type == "moonshot":
		return _count_moonshot_tokens(text, model)
	else:
		# 对于其他提供商，使用基本估算
		return _estimate_tokens(text)


def _count_openai_tokens(text: str, model: Optional[str] = None) -> int:
	"""
	使用tiktoken计算OpenAI模型的token数量
	
	Args:
		text: 要计算的文本
		model: 模型名称，可选
		
	Returns:
		token数量
	"""
	try:
		import tiktoken
		
		# 如果未指定模型，使用默认模型
		encoding_name = "cl100k_base"  # 默认为GPT-4和更新模型的编码
		
		# 根据模型选择编码
		if model:
			model = model.lower()
			if "gpt-3.5" in model:
				encoding_name = "cl100k_base"  # GPT-3.5-Turbo的编码
			elif "gpt-4" in model:
				encoding_name = "cl100k_base"  # GPT-4的编码
			elif "davinci" in model:
				encoding_name = "p50k_base"    # 旧款Davinci模型的编码
		
		# 获取编码器
		try:
			encoding = tiktoken.get_encoding(encoding_name)
		except KeyError:
			logger.warning(f"未知的编码名称: {encoding_name}，使用默认编码")
			encoding = tiktoken.get_encoding("cl100k_base")
			
		# 编码文本并计算token数量
		return len(encoding.encode(text))
		
	except ImportError:
		if "tiktoken" not in _warned_missing_libs:
			logger.warning("未安装tiktoken库，使用近似估算计算OpenAI token")
			_warned_missing_libs.add("tiktoken")
		return _estimate_tokens(text)


def _count_anthropic_tokens(text: str, model: Optional[str] = None) -> int:
	"""
	使用anthropic库计算Anthropic模型的token数量
	
	Args:
		text: 要计算的文本
		model: 模型名称，可选
		
	Returns:
		token数量
	"""
	try:
		from anthropic import Anthropic
		
		# 初始化客户端
		client = Anthropic(api_key="dummy_key")  # 仅用于计算token，不实际调用API
		
		# 计算token数量
		return client.count_tokens(text)
		
	except (ImportError, AttributeError):
		if "anthropic" not in _warned_missing_libs:
			logger.warning("未安装anthropic库或版本过低，使用近似估算计算Claude token")
			_warned_missing_libs.add("anthropic")
		return _estimate_tokens(text)


def _count_moonshot_tokens(text: str, model: Optional[str] = None) -> int:
	"""
	计算Moonshot模型的token数量
	Moonshot与OpenAI使用类似的编码
	
	Args:
		text: 要计算的文本
		model: 模型名称，可选
		
	Returns:
		token数量
	"""
	return _count_openai_tokens(text, "gpt-4")  # 使用GPT-4编码


def _estimate_tokens(text: str) -> int:
	"""
	使用基本方法估算token数量
	根据经验，英文约4个字符一个token，中文约2个字符一个token
	
	Args:
		text: 要计算的文本
		
	Returns:
		估算的token数量
	"""
	if not text:
		return 0
		
	# 计算中文字符数量
	chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
	
	# 计算非中文字符数量
	non_chinese_chars = len(text) - chinese_chars
	
	# 估算token数量
	estimated_tokens = chinese_chars // 2 + non_chinese_chars // 4
	
	# 添加一些额外token，确保不低估
	return max(1, estimated_tokens + len(text) // 100 + 1)


def extract_tokens_from_response(response: Dict[str, Any], provider_type: str) -> Tuple[int, int, int, int]:
	"""
	从API响应中提取token使用信息
	
	Args:
		response: API响应数据
		provider_type: 提供商类型
		
	Returns:
		(输入token数, 输出token数, 缓存输入token数, 缓存输出token数)元组
	"""
	input_tokens = 0
	output_tokens = 0
	cached_input_tokens = 0
	cached_output_tokens = 0
	
	if not isinstance(response, dict):
		return input_tokens, output_tokens, cached_input_tokens, cached_output_tokens
	
	# OpenAI格式
	if provider_type in ["openai", "moonshot"] or (provider_type == "openrouter" and "usage" in response):
		if "usage" in response:
			usage = response["usage"]
			input_tokens = usage.get("prompt_tokens", 0)
			output_tokens = usage.get("completion_tokens", 0)
			
			# 缓存token（仅OpenRouter支持）
			cached_input_tokens = usage.get("cached_prompt_tokens", 0)
			cached_output_tokens = usage.get("cached_completion_tokens", 0)
	
	# Anthropic格式
	elif provider_type == "anthropic" or (provider_type == "openrouter" and "anthropic" in str(response)):
		if "message" in response and "usage" in response.get("message", {}):
			usage = response["message"]["usage"]
			input_tokens = usage.get("input_tokens", 0)
			output_tokens = usage.get("output_tokens", 0)
	
	return input_tokens, output_tokens, cached_input_tokens, cached_output_tokens


def format_token_count(tokens: int) -> str:
	"""
	格式化token数量，使其更易读
	
	Args:
		tokens: token数量
		
	Returns:
		格式化后的字符串
	"""
	if tokens < 1000:
		return str(tokens)
	elif tokens < 1000000:
		return f"{tokens/1000:.1f}K"
	else:
		return f"{tokens/1000000:.1f}M"
